import SimpleITK as sitk
import os
import matplotlib.pyplot as plt
import numpy as np

DCM_DIR = './DcmFiles'
IMAGE_DIR = './ImageFiles'


def calculateWindow(cdf, nda, width, maxCdf=0.99, minCdf=0.1):
    windowMinimum = 0
    windowMaximum = 0

    for i in range(cdf.shape[0] - 1, 0, -1):
        if cdf[i] <= maxCdf:
            windowMaximum = (i * width + nda.min()) // 1
            index = i

            for j in range(index - 1, -1, -1):
                if cdf[j] <= minCdf or j == 0:
                    windowMinimum = (j * width + nda.min()) // 1
                    break

            break
    return windowMinimum, windowMaximum


def dcm2Jpg(dcmPath, imagePath):
    image = sitk.ReadImage(dcmPath)

    nda = sitk.GetArrayFromImage(image)
    nda = nda[0, :, :].reshape(-1)

    # 生成直方图分布
    hist, bin_edges = np.histogram(nda, 1000)

    # width: 区间宽度
    width = bin_edges[1] - bin_edges[0]

    # 求 cdf
    cdf = np.cumsum(hist / sum(hist))

    windowMinimum, windowMaximum = calculateWindow(cdf, nda, width, maxCdf=0.99, minCdf=0.1)

    imgR_255 = sitk.Cast(sitk.IntensityWindowing(image, windowMinimum=windowMinimum,
                                                 windowMaximum=windowMaximum,
                                                 outputMinimum=0.0, outputMaximum=255.0), sitk.sitkUInt8)

    windowMinimum, windowMaximum = calculateWindow(cdf, nda, width, maxCdf=0.992, minCdf=0.08)

    imgG_255 = sitk.Cast(sitk.IntensityWindowing(image, windowMinimum=windowMinimum,
                                                 windowMaximum=windowMaximum + 100,
                                                 outputMinimum=0.0, outputMaximum=255.0), sitk.sitkUInt8)

    windowMinimum, windowMaximum = calculateWindow(cdf, nda, width, maxCdf=0.995, minCdf=0.06)

    imgB_255 = sitk.Cast(sitk.IntensityWindowing(image, windowMinimum=windowMinimum,
                                                 windowMaximum=windowMaximum + 200,
                                                 outputMinimum=0.0, outputMaximum=255.0), sitk.sitkUInt8)

    imgRgb255 = sitk.Compose(imgR_255, imgG_255, imgB_255)

    plt.figure(figsize=(16, 5))
    plt.plot(bin_edges[1:], cdf, color='#000000')
    plt.xlim([windowMinimum, windowMaximum])
    plt.ylim([0, 1])
    plt.grid()
    plt.show()

    # show dcm jpg
    plt.figure(figsize=(10, 10))
    plt.imshow(sitk.GetArrayFromImage(imgRgb255)[0, :, :], cmap=plt.cm.Greys_r)
    plt.axis('off')
    plt.show()

    # dcm convert to jpg
    sitk.WriteImage(sitk.RescaleIntensity(imgRgb255[:, :, 0]), imagePath)
    # sitk.WriteImage(sitk.Cast(sitk.RescaleIntensity(img_blended[:, :, 0]), sitk.sitkUInt8), output_jpg_path)


if __name__ == '__main__':
    dcmFileList = os.listdir(DCM_DIR)

    if '.ipynb_checkpoints' in dcmFileList:
        dcmFileList.remove('.ipynb_checkpoints')

    for dcmFileName in dcmFileList:
        dcmPath = os.path.join(DCM_DIR, dcmFileName)

        imageFileName = dcmFileName.split('.')[0] + '.jpg'
        imagePath = os.path.join(IMAGE_DIR, imageFileName)

        dcm2Jpg(dcmPath, imagePath)
